/**
 * @file population.c
 * @brief Defines interface for a group of modules used in a genetic algorithm
 * @author Scott Moyers
 * @date 2009
 */

#include "population.h"
#include "util.h"

#define DEFAULT_CKT_SIZE        ( 50) /**< default population circuit size */
#define DEFAULT_MAX_MUTATIONS   ( 10) /**< default maximum number of mutations in offspring */
#define DEFAULT_CROSSOVER_PROB  (0.7) /**< default probability of crossover taking place */
/* KEEP_RATIO + BREED_RATIO must be <= 1.0 */
#define DEFAULT_KEEP_RATIO       (0.3) /**< amount of population to keep */
#define DEFAULT_BREED_RATIO      (0.5) /**< amount to be generated from kept ones */

struct Population *new_population()
{
	struct Population *p = (struct Population *) smalloc(sizeof(struct Population));
	init_population(p);
	return p;
}

void init_population(struct Population *p)
{
	p->modules = NULL;
	p->size = 0;
	p->truth_table = NULL;
	p->max_ckt_size = DEFAULT_CKT_SIZE;

	p->selection_funcs = NULL;
	p->n_selection_funcs = 0;

	p->mutation_funcs = NULL;
	p->n_mutation_funcs = 0;
	p->max_mutations = DEFAULT_MAX_MUTATIONS;

	p->crossover_funcs = NULL;
	p->n_crossover_funcs = 0;
	p->crossover_prob = DEFAULT_CROSSOVER_PROB;

	p->keep_ratio = DEFAULT_KEEP_RATIO;
	p->breed_ratio = DEFAULT_BREED_RATIO;

	p->generation = 1;
}

void delete_population(struct Population *p)
{
	if (p != NULL) {
		kill_population(p);
		if (p->truth_table) {
			free_truth_table(p->truth_table);
		}
		free(p->selection_funcs);
		free(p->mutation_funcs);
		free(p->crossover_funcs);
		free(p);
	}
}

void kill_population(struct Population *p)
{
	unsigned int i;
	for(i = 0; i < p->size; i++) {
		free_module(p->modules[i]);
	}
	free(p->modules);
	p->modules = NULL;
}

void generate_population(struct Population *p, dsfmt_t *dsfmt, unsigned int n)
{
	unsigned int i;
	assert(p->truth_table);
	p->modules = (struct Module **) smalloc(sizeof(struct Module *) * n);
	for(i = 0; i < n; i++) {
		p->modules[i] = rand_module(dsfmt, next_module_id(),
					    p->truth_table->nins,
					    p->max_ckt_size);
	}
	p->size = n;
}

void add_selection_func(struct Population *p, SelectionFunc f)
{
	p->selection_funcs = (SelectionFunc *)
		realloc(p->selection_funcs,
			sizeof(SelectionFunc) * (p->n_selection_funcs + 1));
	p->selection_funcs[p->n_selection_funcs] = f;
	p->n_selection_funcs++;
}

void add_selection_funcs(struct Population *p, SelectionFunc f[], unsigned int n)
{
	unsigned int i;
	p->selection_funcs = (SelectionFunc *)
		realloc(p->selection_funcs,
			sizeof(SelectionFunc) * (p->n_selection_funcs + n));
	for(i = 0; i < n; i++) {
		p->selection_funcs[p->n_selection_funcs] = f[i];
		p->n_selection_funcs++;
	}
}

void add_mutation_func(struct Population *p, MutationFunc f)
{
	p->mutation_funcs = (MutationFunc *)
		realloc(p->mutation_funcs,
			sizeof(MutationFunc) * (p->n_mutation_funcs + 1));
	p->mutation_funcs[p->n_mutation_funcs] = f;
	p->n_mutation_funcs++;
}

void add_mutation_funcs(struct Population *p, MutationFunc f[], unsigned int n)
{
	unsigned int i;
	p->mutation_funcs = (MutationFunc *)
		realloc(p->mutation_funcs,
			sizeof(MutationFunc) * (p->n_mutation_funcs + n));
	for(i = 0; i < n; i++) {
		p->mutation_funcs[p->n_mutation_funcs] = f[i];
		p->n_mutation_funcs++;
	}
}

void add_crossover_func(struct Population *p, CrossoverFunc f)
{
	p->crossover_funcs = (CrossoverFunc *)
		realloc(p->crossover_funcs,
			sizeof(CrossoverFunc) * (p->n_crossover_funcs + 1));
	p->crossover_funcs[p->n_crossover_funcs] = f;
	p->n_crossover_funcs++;
}

void add_crossover_funcs(struct Population *p, CrossoverFunc f[], unsigned int n)
{
	unsigned int i;
	p->crossover_funcs = (CrossoverFunc *)
		realloc(p->crossover_funcs,
			sizeof(CrossoverFunc) * (p->n_crossover_funcs + n));
	for(i = 0; i < n; i++) {
		p->crossover_funcs[p->n_crossover_funcs] = f[i];
		p->n_crossover_funcs++;
	}
}

/***************************
 *   Fitness Functions
 ***************************/
#define FIT_FUNCS (3.0f)
#define FIT_SCALE (1.0f/FIT_FUNCS)
static double longest_path_fitness(struct Module *m)
{
	return 1.0f / (depth_module(m) + 1.0f);
}

/**
 * helper function for evaluate_population
 * calculates how correct the circuit performs
 * in comparison the the truth table
 */
static double functionality_fitness(struct Module *m, struct TruthTable *t)
{
	int wrong;
	double fitness;
	execute_module(m);
	wrong = min_signal_diff(m, t);
	fitness = 1.0 - (double)wrong / ((1 << t->nins) * t->nouts);
	return fitness;
}

static double gate_fitness(struct Module *m)
{
	unsigned int i;
	double tot = 1.0;
	for(i = 0; i < m->num_gates; i++) {
		tot += get_gate_weight(&m->circuit[i]);
	}

	return 1.0 / tot;
}

static void evaluate_module_fitness(struct Module *m, struct TruthTable *t)
{
	m->fitness = FIT_SCALE * functionality_fitness(m, t);
	if (m->fitness >= FIT_SCALE) {
		m->fitness += FIT_SCALE * gate_fitness(m);
		m->fitness += FIT_SCALE * longest_path_fitness(m);
	}
}

void evaluate_population(struct Population *p)
{
	unsigned int i;
	for(i = 0; i < p->size; i++) {
		evaluate_module_fitness(p->modules[i], p->truth_table);
	}
}

/**
 * makes a child given 1 parents and the current population
 * @param dsfmt random number generator
 * @param p the population containing mutation and crossover functions
 * @param p1 parent number 1
 * @param p2 parent number 2
 * @return a new child
 */
static struct Module *generate_child(dsfmt_t *dsfmt, struct Population *p,
				     struct Module const *p1, struct Module const *p2)
{
	unsigned int r, j;
	struct Module *child;
	assert(p1->num_gates == p2->num_gates);

	/* do we crossover? */
	if (p->crossover_funcs != 0 && dsfmt_genrand_close_open(dsfmt) < p->crossover_prob) {
		r = dsfmt_genrand_close_open(dsfmt) * p->n_crossover_funcs;
		child = (*p->crossover_funcs[r])(dsfmt, p1, p2);
		assert(child != NULL);
		r = dsfmt_genrand_close_open(dsfmt) * (p->max_mutations + 1);
	} else {
		child = copy_module(p1);
		r = (dsfmt_genrand_close_open(dsfmt) * p->max_mutations) + 1;
	}

	/* mutate */
	if (p->n_mutation_funcs != 0) {
		for(j = 0; j < r; j++) {
			unsigned int rn;
			do {
				rn = dsfmt_genrand_close_open(dsfmt) * p->n_mutation_funcs;
			} while((*p->mutation_funcs[rn])(dsfmt, child));
		}
	}
	return child;
}

/* initial evaluate should have been called before running this */
void next_generation(struct Population *p, dsfmt_t *dsfmt)
{
	unsigned int i;
	int r;
	unsigned int kept = p->size * p->keep_ratio;
	unsigned int children = p->size * p->breed_ratio;
	assert(p->size != 0);

	/* select parents for breeding */
	if (p->n_selection_funcs != 0) {
		r = dsfmt_genrand_close_open(dsfmt) * p->n_selection_funcs;
		(*p->selection_funcs[r])(dsfmt, p->modules, p->size, kept);
	}

	/* generate new children from the kept parents */
	for(i = kept; i < kept + children; i++) {
		/* select 2 parents */
		struct Module const *p1 = p->modules[(int)(dsfmt_genrand_close_open(dsfmt) * kept)];
		struct Module const *p2 = p->modules[(int)(dsfmt_genrand_close_open(dsfmt) * kept)];
		struct Module *child = generate_child(dsfmt, p, p1, p2);

		assert(child != NULL);
		free_module(p->modules[i]);
		p->modules[i] = child;
	}

	/* the rest can be random generations */
	for(i = kept + children; i < p->size; i++) {
		free_module(p->modules[i]);
		p->modules[i] = rand_module(dsfmt, next_module_id(), p->truth_table->nins, p->max_ckt_size);
	}

	/* evaluate new modules */
	for(i = kept; i < p->size; i++) {
		evaluate_module_fitness(p->modules[i], p->truth_table);
	}

	p->generation++;
}

struct Module *get_fittest(struct Population *p)
{
	double fit = -1.0;
	struct Module *fittest = NULL;
	unsigned int i;
	for(i = 0; i < p->size; i++) {
		struct Module *m = p->modules[i];
		if (m->fitness > fit) {
			fittest = m;
			fit = m->fitness;
		}
	}
	return fittest;
}

unsigned int get_fittest_pos(struct Population *p)
{
	double fit = -1.0;
	int fittest = 0;
	unsigned int i;
	for(i = 0; i < p->size; i++) {
		struct Module *m = p->modules[i];
		if (m->fitness > fit) {
			fittest = i;
			fit = m->fitness;
		}
	}
	return fittest;
}

double sum_fitness(struct Population *p)
{
	unsigned int i;
	double sum = 0.0;
	for(i = 0; i < p->size; i++) {
		sum += p->modules[i]->fitness;
	}
	return sum;
}

double average_fitness(struct Population *p)
{
	return sum_fitness(p) / p->size;
}
